package com.example.framework.common.util.thread;

import cn.hutool.core.exceptions.ExceptionUtil;
import cn.hutool.core.thread.ThreadUtil;
import com.alibaba.ttl.threadpool.TtlExecutors;
import com.example.framework.common.constant.BaseConstants;
import com.example.framework.common.util.TraceIdGenerateUtils;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.*;

@Slf4j
public class TTLTraceThreadPoolExecutor extends ThreadPoolTaskExecutor {

    @Value("${spring.application.name:}")
    @Getter
    private String applicationName;

    private final ExecutorService ttlExecutorService;

    public TTLTraceThreadPoolExecutor(ExecutorService executorService) {
        this.ttlExecutorService = TtlExecutors.getTtlExecutorService(executorService);
    }

    public TTLTraceThreadPoolExecutor(int threadCount, int maxThreadCount, int queueSize, String threadName, RejectedExecutionHandler rejectedExecutionHandler) {
        ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(
                threadCount,
                maxThreadCount,
                120,
                TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(queueSize),
                ThreadUtil.createThreadFactory(threadName + "-"), rejectedExecutionHandler);
        this.ttlExecutorService = TtlExecutors.getTtlExecutorService(threadPoolExecutor);
    }

    public TTLTraceThreadPoolExecutor(int threadCount, int maxThreadCount, int queueSize, int keepAliveSeconds, String threadName, RejectedExecutionHandler rejectedExecutionHandler) {
        ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(
                threadCount,
                maxThreadCount,
                keepAliveSeconds,
                TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(queueSize),
                ThreadUtil.createThreadFactory(threadName + "-"), rejectedExecutionHandler);
        this.ttlExecutorService = TtlExecutors.getTtlExecutorService(threadPoolExecutor);
    }

    public void execute(Runnable task) {
        String curTraceId = getCurTraceId(); // 获取当前线程的traceId
        ttlExecutorService.execute(() -> {
            String preTraceId = MDC.get(BaseConstants.TRACE_UNIQUE_ID); // 获取子线程的traceId
            // 在子线程中设置traceId
            MDC.put(BaseConstants.TRACE_UNIQUE_ID, curTraceId);
            try {
                task.run();
            } catch (Exception e) {
                log.error("submit error {}", ExceptionUtil.getMessage(e));
            } finally {
                if (preTraceId != null) { // 恢复现场
                    MDC.put(BaseConstants.TRACE_UNIQUE_ID, preTraceId);
                } else {
                    // 清理traceId
                    MDC.remove(BaseConstants.TRACE_UNIQUE_ID);
                }
            }
        });
    }

    /**
     * 提交任务
     * submit方法返回一个Future对象，可以通过Future的get方法获取任务执行结果
     */
    public <T> Future<T> submit(Callable<T> task) {
        String curTraceId = getCurTraceId(); // 获取当前线程的traceId
        return ttlExecutorService.submit(() -> {
            String preTraceId = MDC.get(BaseConstants.TRACE_UNIQUE_ID); // 获取子线程的traceId
            // 在子线程中设置traceId
            MDC.put(BaseConstants.TRACE_UNIQUE_ID, curTraceId);
            try {
                return task.call();
            } catch (Exception e) {
                log.error("submit error {}", ExceptionUtil.getMessage(e));
            } finally {
                if (preTraceId != null) { // 恢复现场
                    MDC.put(BaseConstants.TRACE_UNIQUE_ID, preTraceId);
                } else {
                    // 清理traceId
                    MDC.remove(BaseConstants.TRACE_UNIQUE_ID);
                }
            }
            return null;
        });
    }

    /**
     * 获取父线程的traceId
     */
    private String getCurTraceId() {
        String traceId = MDC.get(BaseConstants.TRACE_UNIQUE_ID);
        if (traceId == null) {
            // 生成一个traceId
            traceId = TraceIdGenerateUtils.traceId(applicationName);
            MDC.put(BaseConstants.TRACE_UNIQUE_ID, traceId);
        }
        return traceId;
    }

}
